/* Copyright (c) 2024 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 * ===================================================================================================================*/

#include "external/register/register_custom_pass.h"
#include "register/custom_pass_helper.h"
#include "graph/debug/ge_log.h"
#include "graph/debug/ge_util.h"
#include "common/plugin/plugin_manager.h"
#include "register/custom_pass_context_impl.h"

namespace ge {
PassReceiver::PassReceiver(PassRegistrationData &reg_data) {
  CustomPassHelper::Instance().Insert(reg_data);
}

class PassRegistrationDataImpl {
 public:
  PassRegistrationDataImpl() = default;
  ~PassRegistrationDataImpl() = default;

  explicit PassRegistrationDataImpl(const std::string &pass_name);

private:
  friend class PassRegistrationData;
  std::string pass_name_;
  CustomPassFunc custom_pass_fn_ = nullptr;
};

PassRegistrationDataImpl::PassRegistrationDataImpl(const std::string &pass_name)
    : pass_name_(pass_name),
      custom_pass_fn_(nullptr) {}

PassRegistrationData::PassRegistrationData(std::string pass_name) {
  impl_ = ge::ComGraphMakeShared<PassRegistrationDataImpl>(pass_name);
  if (impl_ == nullptr) {
    GELOGW("[Check][Param] make impl failed, pass_name:%s", pass_name.c_str());
  }
}

std::string PassRegistrationData::GetPassName() const {
  if (impl_ == nullptr) {
    return "";
  }
  return impl_->pass_name_;
}

PassRegistrationData &PassRegistrationData::CustomPassFn(const CustomPassFunc &custom_pass_fn) {
  if (impl_ != nullptr) {
    impl_->custom_pass_fn_ = custom_pass_fn;
  }
  return *this;
}

CustomPassFunc PassRegistrationData::GetCustomPassFn() const {
  if (impl_ == nullptr) {
    return nullptr;
  }
  return impl_->custom_pass_fn_;
}

CustomPassContext::CustomPassContext() {
  impl_ = ComGraphMakeUnique<CustomPassContextImpl>();
  if (impl_ == nullptr) {
    GELOGW("[Check][Param] make impl failed");
  }
}

void CustomPassContext::SetErrorMessage(const AscendString &error_message) {
  if (impl_ != nullptr) {
    impl_->SetErrorMessage(error_message);
  }
}

AscendString CustomPassContext::GetErrorMessage() const {
  if (impl_ != nullptr) {
    return impl_->GetErrorMessage();
  }
  return "";
}

CustomPassHelper &CustomPassHelper::Instance() {
  static CustomPassHelper instance;
  return instance;
}

void CustomPassHelper::Insert(const PassRegistrationData &reg_data) {
  (void)registration_datas_.emplace_back(reg_data);
}

Status CustomPassHelper::Load() {
  GELOGD("[Load][CustomPassLibs] Start to load custom pass libs");
  std::string opp_path;
  ge::PluginManager::GetOppPath(opp_path);
  std::string vendors_path = opp_path + "/vendors";

  // 存储所有的 .so 文件路径
  std::vector<std::string> so_files;
  ge::PluginManager::FindSoFilesInCustomPassDirs(vendors_path, so_files);

  if (so_files.empty()) {
    GELOGD("No custom pass libs found in %s, skip loading custom pass libs", vendors_path.c_str());
    return ge::SUCCESS;
  }

  // 逐个 dlopen
  for (const auto &so_file : so_files) {
    void *handle = dlopen(so_file.c_str(), RTLD_NOW | RTLD_LOCAL);
    if (!handle) {
      const char* error = dlerror();
      REPORT_INPUT_ERROR("E19029", std::vector<std::string>({"passlibname", "reason"}),
                         std::vector<std::string>({so_file, error}));
      GELOGE(ge::FAILED, "Failed to load %s: %s", so_file.c_str(), error);
      return ge::FAILED;
    }
    handles_.emplace_back(handle);
    GELOGI("Load custom pass lib %s success", so_file.c_str());
  }
  return ge::SUCCESS;
}

Status CustomPassHelper::Unload() {
  registration_datas_.clear();
  for (auto &handle : handles_) {
    if (handle != nullptr && dlclose(handle) != 0) {
      GELOGE(ge::FAILED, "[Unload][CustomPassLibs] Failed to unload custom pass lib: %s", dlerror());
      return ge::FAILED;
    }
    GELOGI("Unload custom pass lib success");
  }
  handles_.clear();
  return ge::SUCCESS;
}

Status CustomPassHelper::Run(ge::GraphPtr &graph, CustomPassContext &custom_pass_context) const {
  for (auto &item : registration_datas_) {
    GELOGD("Start to run custom pass [%s]!", item.GetPassName().c_str());
    const auto custom_pass_fn = item.GetCustomPassFn();
    if (custom_pass_fn == nullptr) {
      GELOGW("[Check][Param] Get custom_pass_fn of custom pass %s failed", item.GetPassName().c_str());
      continue;
    }
    const auto ret = custom_pass_fn(graph, custom_pass_context);
    if (ret != SUCCESS) {
      GE_LOGE("Custom pass [%s] run failed!", item.GetPassName().c_str());
      REPORT_INPUT_ERROR("E19028", std::vector<std::string>({"passname", "retcode", "reason"}),
                         std::vector<std::string>({item.GetPassName(), std::to_string(ret),
                         std::string(custom_pass_context.GetErrorMessage().GetString())}));
      return FAILED;
    }
    GELOGD("Run custom pass [%s] success!", item.GetPassName().c_str());
  }
  return SUCCESS;
}
}  // namespace ge
